import  h5py
import os
import SimpleITK as sitk
# import nibabel as nib
from tqdm import tqdm
import numpy as np
from multiprocessing import Process

def rescale_intensity(volume, percentils=[0.5, 99.5], bins_num=256, norm=False):
    obj_volume = volume[np.where(volume > 0)]
    min_value = np.percentile(obj_volume, percentils[0])
    max_value = np.percentile(obj_volume, percentils[1])

    if bins_num == 0:
        obj_volume = (obj_volume - min_value) / (max_value - min_value).astype(np.float32)
    else:
        obj_volume = np.round((obj_volume - min_value) / (max_value - min_value) * (bins_num - 1))
        obj_volume[np.where(obj_volume < 1)] = 1
        obj_volume[np.where(obj_volume > (bins_num - 1))] = bins_num - 1

    volume = volume.astype(obj_volume.dtype)
    volume[np.where(volume > 0)] = obj_volume
    if norm:
        volume = volume.astype(float) / (bins_num - 1)

    return volume

def clipseScaleSitkImage(sitk_image,low=5, up=95):
    np_image = sitk.GetArrayFromImage(sitk_image)
    # threshold image between p10 and p98 then re-scale [0-255]
    p0 = np_image.min().astype('float')
    p10 = np.percentile(np_image, low)
    p99 = np.percentile(np_image, up)
    p100 = np_image.max().astype('float')
    # logger.info('p0 {} , p5 {} , p10 {} , p90 {} , p98 {} , p100 {}'.format(p0,p5,p10,p90,p98,p100))
    sitk_image = sitk.Threshold(sitk_image,
                                lower=p10,
                                upper=p100,
                                outsideValue=p10)
    sitk_image = sitk.Threshold(sitk_image,
                                lower=p0,
                                upper=p99,
                                outsideValue=p99)
    sitk_image = sitk.RescaleIntensity(sitk_image,
                                       outputMinimum=0,
                                       outputMaximum=255)
    return sitk_image

def sitkResize(image, new_size, interpolator):
    resample = sitk.ResampleImageFilter()
    resample.SetInterpolator(interpolator)
    resample.SetOutputDirection(image.GetDirection())
    resample.SetOutputOrigin(image.GetOrigin())
    new_spacing = [sz * spc / nsz for nsz, sz, spc in zip(new_size, image.GetSize(), image.GetSpacing())]
    resample.SetOutputSpacing(new_spacing)
    orig_size = np.array(image.GetSize(), dtype=np.int32)
    orig_spacing = list(image.GetSpacing())
    new_size=[oz*os/nz for oz,os,nz in zip(orig_size,orig_spacing,new_spacing)]
    new_size = np.ceil(new_size).astype(np.int32)  # Image dimensions are in integers
    new_size = [int(s) for s in new_size]
    resample.SetSize(new_size)
    newimage = resample.Execute(image)
    return newimage

from util.dir import mkdir_if_not_exist
import glob
def convert_h5(subject, data_path, out_path,slice_start=0.2,slice_end=0.8):
    # h5_file = os.path.join(out_path, subject+'.h5')
    # h5_file = h5py.File(h5_file, 'w')
    modalities = ['Mag_Images', 'CBV_REG', 'Pha_Images','mIP_Images(SW)','SWI_Images']
    for modality in modalities:
        img = glob.glob(f"{data_path}/{subject}/*{modality}*.nii")
        assert len(img) == 1,print(f"{data_path}/{subject}/*{modality}")
        img=img[0]
        if os.path.exists(img):
            img_nii = sitk.ReadImage(img)
            img_nii=sitkResize(img_nii,[256,256,img_nii.GetSize()[-1]],sitk.sitkLinear)
            img_nii=clipseScaleSitkImage(img_nii,0,100)
            img_nii = img_nii[...,int(img_nii.GetSize()[-1]* slice_start):int(img_nii.GetSize()[-1]* slice_end)]
            assert (img_nii.GetSize()[0] == 256 and img_nii.GetSize()[1] == 256)
            mkdir_if_not_exist(f"{out_path}/{subject}/")
            sitk.WriteImage(img_nii,f"{out_path}/{subject}/{modality}.nii.gz")

            # h5_file[modality] = img_data
    # h5_file.close()

data_path = r'../../../data/DSC_DWI_SWI/'
out_path = r'./data_resize'
os.makedirs(out_path, exist_ok=True)
all_subject = os.listdir(data_path)

num_processes = 5

for loop in tqdm(range(len(all_subject) // num_processes + 1)):
    processes = []
    for subject in all_subject[loop*num_processes:(loop+1)*num_processes]:
        p = Process(target=convert_h5, args=(subject,data_path, out_path))
        processes.append(p)
        p.start()
    for p in processes:
        p.join()
